#define TILE_WIDTH 32
#define ROWS 50    // C的行数和A的行数
#define COLS 1604  // A的列数和B的行数

__global__ void matrixMultiplyKernel(cuFloatComplex* A, cuFloatComplex* B, cuFloatComplex* C, int numMatrices) {
    // 获取矩阵对索引和C的行列索引
    int p = blockIdx.x; // 第p个矩阵对
    if (p >= numMatrices) return;

    int row = blockIdx.y * TILE_WIDTH + threadIdx.y; // C的行索引
    int col = blockIdx.z * TILE_WIDTH + threadIdx.x; // C的列索引

    // 检查是否超出C的边界（50x50）
    if (row < ROWS && col < ROWS) {
        // 分配共享内存
        __shared__ cuFloatComplex A_shared[TILE_WIDTH][TILE_WIDTH];
        __shared__ cuFloatComplex B_shared[TILE_WIDTH][TILE_WIDTH];

        // 初始化结果
        cuFloatComplex sum = make_cuFloatComplex(0.0f, 0.0f);

        // 计算A和B的基地址
        int base_idx_A = p * ROWS * COLS + row * COLS; // A的行起始地址
        int base_idx_B = p * COLS * ROWS + col;        // B的列起始地址

        // 计算分块次数
        int num_tiles = (COLS + TILE_WIDTH - 1) / TILE_WIDTH; // ceil(1604/32) = 51

        // 分块循环
        for (int t = 0; t < num_tiles; t++) {
            // 计算当前子块的k索引
            int k = t * TILE_WIDTH + threadIdx.x;

            // 加载A的子块到共享内存
            if (threadIdx.y < TILE_WIDTH && k < COLS) {
                A_shared[threadIdx.y][threadIdx.x] = A[base_idx_A + k];
            } else {
                A_shared[threadIdx.y][threadIdx.x] = make_cuFloatComplex(0.0f, 0.0f); // 填充0
            }

            k = t * TILE_WIDTH + threadIdx.y;
            // 加载B的子块到共享内存
            if (threadIdx.x < TILE_WIDTH && k < COLS) {
                B_shared[threadIdx.y][threadIdx.x] = B[base_idx_B + k * ROWS];
            } else {
                B_shared[threadIdx.y][threadIdx.x] = make_cuFloatComplex(0.0f, 0.0f); // 填充0
            }

            // 同步，确保共享内存加载完成
            __syncthreads();

            // 计算部分点积
            for (int m = 0; m < TILE_WIDTH; m++) {
                sum = cuCaddf(sum, cuCmulf(A_shared[threadIdx.y][m], B_shared[m][threadIdx.x]));
            }

            // 同步，准备下一次加载
            __syncthreads();
        }

        // 写入结果到C
        C[p * ROWS * ROWS + row * ROWS + col] = sum;
    }
}